#include "header.h"

static fftwq_plan plan_f1, plan_b1;
static fftwq_plan plan_fr1, plan_br1, plan_br2;
static __float128 *k, *u, overN, q, *mk, dtau, gama, sp, num, den, err, err2, l2, tol, hmax, ml0, ml, htt;
static __float128 coeff, aa, bb, NV, *ra[4], *rt1, *rt2, *rhu[1], *rhk[2], *redk[2], *reDk, *reD;
static __float128 ip, *reDk, *reRk;
static int N, n2, nthreads; 
static work_ptr tmp;
static pthread_t thread1, thread2;
static pthread_mutex_t mut;
static pthread_cond_t cvar, cvar2;

void err_msg(char* str)
{
  printf("%s\n",str);
  exit(1);
}

__float128 signum(__float128 x) {
  if (x > 0.Q) return  1.Q;
  if (x < 0.Q) return -1.Q;
  return 0.Q;
}

__float128 signum2(__float128 x) {
  if (x > 0.Q) return  1.Q;
  else return -1.Q;
}

void init(int *argc, char **argv, params_ptr in, work_ptr wrk)
{
  char str[480], line[480], value[480], param[480];
  sprintf(str, "Usage:\n\t%s input", argv[0]);
  if (*argc != 2) err_msg(str);

  sprintf(str, "%s.cfg", argv[1]);
  FILE *fh = fopen(str,"r"); 
  FILE *fhlog = fopen("run.log","w"); 
  if (fh == NULL) err_msg("Cannot open file");
  while (fgets(line, 480, fh)!=NULL) {
    sscanf(line, "%s\t%s", param, value);
    if (strcmp(param,"#runname=") == 0) sprintf(in->runname,"%s", value);
    if (strcmp(param,"#resname=") == 0) sprintf(in->resname,"%s", value);
    if (strcmp(param,"#npoints=") == 0) in->n = atoi(value);
    if (strcmp(param,"#gravity=") == 0) in->g = strtoflt128 (value, NULL);
    if (strcmp(param,"#velocit=") == 0) in->c = strtoflt128 (value, NULL);
    if (strcmp(param,"#delta_t=") == 0) in->dtau = strtoflt128 (value, NULL);
    if (strcmp(param,"#toleran=") == 0) in->tol = strtoflt128 (value, NULL);
    if (strcmp(param,"#perturb=") == 0) in->alfa = strtoflt128 (value, NULL);
    if (strcmp(param,"#io_time=") == 0) in->t_skip = atoi(value);
    if (strcmp(param,"#nthread=") == 0) in->nthreads = atoi(value);
    if (strcmp(param,"#hmax_tt=") == 0) htt = strtoflt128 (value, NULL);
  }
  fprintf(fhlog, "runname = %s\n", in->runname );
  fprintf(fhlog, "resname = %s\n", in->resname );
  fprintf(fhlog, "npoints = %d\n", in->n );
  fprintf(fhlog, "gravity = %Qe\n", in->g );
  fprintf(fhlog, "velocit = %Qe\n", in->c );
  fprintf(fhlog, "toleran = %Qe\n", in->tol );
  fprintf(fhlog, "delta_t = %Qe\n", in->dtau );
  fprintf(fhlog, "perturb = %Qe\n", in->alfa );
  fprintf(fhlog, "hmax_tt = %Qe\n", htt );
  fprintf(fhlog, "io_time = %d\n", in->t_skip );
  fprintf(fhlog, "nthread = %d\n", in->nthreads );
  fclose(fh); fclose(fhlog);
  N = in->n; n2 = N/2 + 1; nthreads = in->nthreads;

  if (nthreads != 1) {
   fftwq_init_threads();
   fftwq_plan_with_nthreads(nthreads);
  }

  ra[0] = fftwq_malloc(sizeof(__float128)*n2);
  ra[1] = fftwq_malloc(sizeof(__float128)*n2);
  ra[2] = fftwq_malloc(sizeof(__float128)*n2);
  ra[3] = fftwq_malloc(sizeof(__float128)*n2);
  wrk->YY  = fftwq_malloc(sizeof(__float128)*n2);
  wrk->YYk = fftwq_malloc(sizeof(__float128)*n2);

  rhu[0] = fftwq_malloc(sizeof(__float128)*n2);
  rhk[0] = fftwq_malloc(sizeof(__float128)*n2);
  rhk[1] = fftwq_malloc(sizeof(__float128)*n2);

  redk[0] = fftwq_malloc(sizeof(__float128)*n2);
  redk[1] = fftwq_malloc(sizeof(__float128)*n2);

  reDk = fftwq_malloc(sizeof(__float128)*n2);
  reD  = fftwq_malloc(sizeof(__float128)*n2);
  reRk  = fftwq_malloc(sizeof(__float128)*n2);

  //u  = fftwq_malloc(sizeof(__float128)*(in->n));
  q  = (in->g)/powq(in->c,2);
  
  sprintf(value,"r2r.fftw.%d", nthreads);
  if (fftwq_import_wisdom_from_filename(value) == 0) {
    printf("Creating plans\n");
    fftwq_set_timelimit(1.);
    plan_fr1 = fftwq_plan_r2r_1d(n2, ra[0], ra[1], FFTW_REDFT00, fmode);
    plan_br1 = fftwq_plan_r2r_1d(n2, ra[1], ra[0], FFTW_REDFT00, fmode);
    plan_br2 = fftwq_plan_r2r_1d(n2, ra[3], ra[2], FFTW_REDFT00, fmode);
    fftwq_export_wisdom_to_filename(value);
    err_msg("Plans created. Rerun the computation");
  } else {
    fftwq_set_timelimit(1.);
    plan_fr1 = fftwq_plan_r2r_1d(n2, ra[0], ra[1], FFTW_REDFT00, fmode);
    plan_br1 = fftwq_plan_r2r_1d(n2, ra[1], ra[0], FFTW_REDFT00, fmode);
    plan_br2 = fftwq_plan_r2r_1d(n2, ra[3], ra[2], FFTW_REDFT00, fmode);
  }
  for (int j = 0; j < in->n; j ++) {
    //u[j] = 2*pi*j/(in->n);
    //wrk->Y[j] = 0.2*exp(-2*pow(u[j]-pi,2)); // -cos(3.*u[j])
  }
  fh = fopen(in->resname,"r");
  int jj = 0;
  if (fh != NULL) {
    char v1[80],v2[80],v3[80],v4[80], *dummy;
    while (fgets(line, 480, fh)!=NULL) {
      if (jj == 0) {
        dummy = fgets(line, 160, fh);
        dummy = fgets(line, 160, fh);
      }
      sscanf(line, "%s\t%s\t%s\t%s\n",  v1, v2, v3, v4);
      if (jj == n2) break;
      wrk->YY[jj] = strtoflt128(v4, NULL);
      jj++;
      
    }
    printf("Lines read = %d of %d expected\n", jj, n2);
    if ( jj != n2) err_msg("I/O error\n");
    fclose(fh);
  }
  

  overN = 1.Q/N; 
  in->w_ptr = wrk; dtau = in->dtau;
  gama = 1.Q + 1.Q/dtau; tol = 1.e-10Q;
  fft_even(wrk->YY,wrk->YYk);
  for(int j = 0; j < n2; j++) wrk->YYk[j] = (wrk->YYk[j])*overN*cexpq((in->alfa)*j); 

  /*if (jj == 0) {
    wrk->Yk[0] = 0;
    for (int j = 1; j < N; j++) wrk->Yk[0] += -cabsq(k[j])*conj(wrk->Yk[j])*(wrk->Yk[j]); 
  }*/
}

void khat_even(__float128 *in, __float128 *out) {
  memcpy(ra[0], in, sizeof(__float128)*n2);
  writer_out(ra[0], "orig_r2r.dat");
  fftwq_execute(plan_fr1);
  writer_out(ra[1], "fft_r2r.dat");
  for (int j = 0; j < n2 ; j++) ra[1][j] = j*overN*ra[1][j];
  fftwq_execute(plan_br1);
  memcpy(out, ra[0], sizeof(__float128)*n2);
}

void fft_even(__float128 *in, __float128 *out) {
  memcpy(ra[0], in, n2*sizeof(__float128));
  fftwq_execute(plan_fr1);
  memcpy(out, ra[1], n2*sizeof(__float128));
}

void ifft_even(__float128 *in, __float128 *out) {
  memcpy(ra[1], in, n2*sizeof(__float128));
  fftwq_execute(plan_br1);
  memcpy(out, ra[0], n2*sizeof(__float128));
}

void ifft_even_two(__float128 *in, __float128 *out) {
  memcpy(ra[3], in, n2*sizeof(__float128));
  fftwq_execute(plan_br2);
  memcpy(out, ra[2], n2*sizeof(__float128));
}

__float128 inner_product_real(__float128 *in1, __float128 *in2) {
  ip = 0.Q;
  for (int j = 0; j < n2; j++) {
    ip += 2.Q*in1[j]*in2[j];
  }
  ip = ip - (in1[0]*in2[0] + in1[n2-1]*in2[n2-1]);
  return ip;
}

__float128 mi_product_real(__float128 *in1, __float128 *in2) {
  ip = 0.Q;
  for (int j = 0; j < n2; j++) {
    ip += 2.Q*in1[j]*in2[j]/(j - q);
  }
  ip = ip - (in1[0]*in2[0]/(-q) + in1[n2-1]*in2[n2-1]/(n2-1-q));
  return ip;
}

void prepare_rhs_even(work_ptr wrk) {
  ifft_even(wrk->YYk,wrk->YY); 
  for (int j = 0; j < n2; j++) {
    rhk[0][j] = (wrk->YYk[j])*j;
    rhk[1][j] = (wrk->YY[j])*(wrk->YY[j]);
  }
  fft_even(rhk[1],rhk[1]);  
  ifft_even(rhk[0],rhu[0]);
  for (int j = 0; j < n2; j++) {
    rhk[1][j] = j*rhk[1][j]*overN;   
    rhk[0][j] = rhu[0][j]*(wrk->YY[j])*overN;
  }
  fft_even(rhk[0],rhk[0]); 
  for (int j = 0; j < n2; j++) {
    rhk[0][j] = -(j - q)*(wrk->YYk[j]) + q*(0.5Q*rhk[1][j] + rhk[0][j]); 
  }
}


void *run_thread_one(void *in) {
  work_ptr WRK  = (work_ptr) in;

  pthread_mutex_lock(&mut);
  ifft_even(reDk, reD);
  pthread_mutex_unlock(&mut);

  for (int j = 0; j < n2; j++) redk[0][j] = (WRK->YY[j])*reD[j]; 

  //pthread_mutex_lock(&mut);
  fft_even(redk[0], redk[0]);
  //pthread_mutex_unlock(&mut);

  /*pthread_cond_wait(&cvar, &mut);
  for (int j = 0; j < (n2-1)/2; j++) {
    redk[0][j] = -(j - q)*reDk[j] + q*(j*overN*redk[0][j] + redk[1][j]);  //half here
  }*/
					
  pthread_exit(NULL);
}

void *run_thread_two(void *in) {
  work_ptr WRK  = (work_ptr) in;
  for (int j = 0; j < n2; j++)  redk[1][j] = j*reDk[j]; 

  //pthread_mutex_lock(&mut);
  ifft_even_two(redk[1], redk[1]);
  //pthread_mutex_unlock(&mut);

  pthread_mutex_lock(&mut);
  for (int j = 0; j < n2; j++) redk[1][j] = overN*((WRK->YY[j])*redk[1][j] + rhu[0][j]*reD[j]);
  ifft_even_two(redk[1], redk[1]);
  pthread_mutex_unlock(&mut); 

  /*pthread_cond_signal(&cvar); 
  for (int j = (n2-1)/2; j < n2; j++) {   
    redk[0][j] = -(j - q)*reDk[j] + q*(j*overN*redk[0][j] + redk[1][j]);    //half here
  }*/
  
					
  pthread_exit(NULL);
}

void prepare_linearization_even(work_ptr wrk) {
  int iret1, iret2;
  void *aux = (void *)wrk;

  pthread_mutex_init(&mut,NULL);
  pthread_cond_init(&cvar,NULL);
  pthread_cond_init(&cvar2,NULL);

  iret1 = pthread_create( &thread1, NULL, &run_thread_one, aux);
  iret2 = pthread_create( &thread2, NULL, &run_thread_two, aux);

  pthread_join(thread1, NULL);
  pthread_join(thread2, NULL);
  //err_msg("Complete");

  /*for (int j = 0; j < n2; j++) {
    redk[1][j] = j*reDk[j]; 
  }   
  ifft_even(reDk, reD);
  ifft_even(redk[1], redk[1]);

  for (int j = 0; j < n2; j++) {
    redk[0][j] = (wrk->YY[j])*reD[j];      
    redk[1][j] = overN*((wrk->YY[j])*redk[1][j] + rhu[0][j]*reD[j]); 
  }
  fft_even(redk[0], redk[0]);
  fft_even(redk[1], redk[1]);*/
  
  for (int j = 0; j < n2; j++) {
    redk[0][j] = -(j - q)*reDk[j] + q*(j*overN*redk[0][j] + redk[1][j]);
  }
}

void iterate_cg(work_ptr wrk) {
  int counter, counter2;
  __float128 cg_error = 1.Q;
  __float128 err = 1.Q;
  __float128 eps = 0.01Q;

  counter2 = 1; NV = 1.Q; counter = 0;
  while (NV > tol) {
    prepare_rhs_even(wrk);
    writer_out(wrk->YY,"restart.dat");  
    writer_out(wrk->YYk,"stokes_spec.dat");
    for (int j = 0; j < n2; j++) {
      rhk[1][j] = 0.Q;
      reRk[j] = -1.Q*rhk[0][j];
      reDk[j] = reRk[j]/(j - q);
    }
    NV = sqrtq(inner_product_real(rhk[0],rhk[0]) );  
    if ((counter2%100) == 0 ) {
        hmax = ((wrk->YY[n2-1])-(wrk->YY[0]))/(2.Q*pi);
	printf("Nevyazka = %Qe Newton cycle %1d\n", NV, counter2);
    }
    counter = 0; cg_error = 1.Q;
    while (cg_error > eps*NV ) {       // start CG cycle 
      prepare_linearization_even(wrk);
      num = mi_product_real(reRk,reRk);    
      den = inner_product_real(reDk,redk[0]);
      aa = num/den;
      den = mi_product_real(reRk, reRk);
      for (int j = 0; j < n2; j++) {
        rhk[1][j] +=  aa*reDk[j]; 
        reRk[j]   += -aa*redk[0][j];
      }
      num = -aa*mi_product_real(redk[0], reRk);  // Polak-Ribiere
      //num = mi_product_real(reRk, reRk[0]);      // Fletcher-Reeves

      bb = num/den;
      for (int j = 0; j < n2; j++) reDk[j] = reRk[j]/(j - q) + bb*reDk[j];
      cg_error = sqrtq(inner_product_real(reRk,reRk));
      counter++;
      if ((counter%5)==0) printf("NV = %Qe CG Error = %Qe at %d step\n", NV, cg_error, counter);
      if (counter == 1000) break;

    }                                  // end of CG cycle
    if (counter == 1000) break;
    for (int j = 0; j < n2; j++) wrk->YYk[j] = wrk->YYk[j] + rhk[1][j];
    hmax = ((wrk->YY[n2-1])-(wrk->YY[0]))/(2.Q*pi);
    printf("NV = %Qe at %3d CG step, H/L = %.12Qe Speed = %.12Qe\n", NV, counter, hmax, 1.Q/sqrtq(q));

    counter2++;
   
  }  // end of outer cycle

  if (counter == 1400) err_msg("CG diverged");

  ifft_even(wrk->YYk,wrk->YY);  
  writer_out(wrk->YY,"stokes_phys.dat");
  //writer_out(wrk->YY,"restart.dat");
  writer_out(wrk->YYk,"stokes_spec.dat");
  FILE *fid = fopen("steep.data","a");
  fprintf(fid, "%.12Qe\t%.12Qe\n", hmax, 1.Q/sqrtq(q));
  fclose(fid);
  //err_msg("Done");
  //pthread_exit(NULL);
}

void *run_cg_thread_one() {


}

void *run_cg_thread_two() {


}

/*void iterate_cr(work_ptr wrk) {
  int counter, counter2;
  __float128 cr_error = 1.Q;
  __float128 err = 1.Q;
  __float128 eps = 0.01Q;

  counter2 = 1; NV = 1.Q; counter = 0;
  while (NV > tol) {
    prepare_rhs_even(wrk);
    NV = sqrtq(inner_product_real(rhk[0],rhk[0]) );  

    for (int j = 0; j < n2; j++) {
      reDyk[j] = 0.Q;			//  x0
      reRk[j] = -1.Q*rhk[0][j]/(j - q); //  r0
    }
    memcpy(reDk, reRk, n2*sizeof(__float128));  // p0
    prepare_linearization_even(wrk);


    if ((counter2%1) == 0 ) {
        hmax = ((wrk->YY[n2-1])-(wrk->YY[0]))/(2.Q*pi);
	printf("Nevyazka = %Qe Newton cycle %1d\n", NV, counter2);
    }
    counter = 0; cr_error = 1.Q;
    while (cr_error > eps*NV ) {       // start CR cycle       

      for (int j = 0; j < n2; j++) rMiRk[j] = Apk[j]/(j - q);   //M^-1 A p

      num = inner_product_real(reRk,redk[0]);
      den = inner_product_real(Apk, rMiRk);
      aa = num/den;				       // alpha_k

      memcpy(rhk[1], reRk, n2*sizeof(__float128));     //  r_k
      memcpy(rhk[0], redk[0], n2*sizeof(__float128));  //  Ar_k

      for (int j = 0; j < n2; j++) {
        reDyk[j]  = reDyk[j] + aa*reDk[j];    //  x_k+1    
        reRk[j]   = reRk[j] - aa*rMiRk[j];    //  r_k+1
      }

      prepare_linearization_even(wrk);        // Ar_k+1

      num = inner_product_real(reRk, redk[0]);
      den = inner_product_real(rhk[1], rhk[0]);
      bb = num/den;			      // beta_k

      for (int j = 0; j < n2; j++) {
	reDk[j] = reRk[j] + bb*reDk[j];
	Apk[j]  = redk[0][j] + bb*Apk[j];
      }
      cr_error = sqrtq(inner_product_real(reRk,reRk));
      counter++;
      if ((counter%1)==0) printf("NV = %Qe CG Error = %Qe at %d step\n", NV, cr_error, counter);
      if (counter == 1000) break;

    }                                  // end of CR cycle
    if (counter == 1000) break;
    for (int j = 0; j < n2; j++) wrk->YYk[j] = wrk->YYk[j] + reDyk[j];
    printf("Error = %Qe at %3d CG step, H/L = %Qe Speed = %.12Qe\n", cr_error, counter, hmax, 1.Q/sqrtq(q));
    counter2++;
   
  }  // end of outer cycle

  if (counter == 1400) err_msg("CR diverged");

  ifft_even(wrk->YYk,wrk->YY);  
  writer_out(wrk->YY,"stokes_phys.dat");
  writer_out(wrk->YY,"restart.dat");
  writer_out(wrk->YYk,"stokes_spec.dat");
  FILE *fid = fopen("steep.data","a");
  fprintf(fid, "%.12Qe\t%.12Qe\n", hmax, 1.Q/sqrtq(q));
  fclose(fid);
  err_msg("Done");
  pthread_exit(NULL);
}*/

/*void iterate(work_ptr wrk) {
  err = 1;
  int nn = 0;
  char st[80];

  write_out(wrk->Yk,"before.txt");
  while (err > tol) { 
    prepare_rhs(wrk); num = 0.; den = 0.;
    for (int j = 0; j < N; j++) {
      num += conj(wrk->Yk[j])*(hk[3][j]);
      den += mk[j]*pow(cabs(wrk->Yk[j]),2);
    }
    sp = 1.*num/den;
    err = 0.; l2 = 0; err2;
    for (int j = 0; j < N; j++) {
      rhs_yk = ((hk[3][j])/mk[j] - gama*sp*(wrk->Yk[j]) )*dtau;
      wrk->Yk[j] += rhs_yk;
      if (j > N/2) wrk->Yk[j] = conj(wrk->Yk[N-j]);
    }
    for (int j = 0; j < N; j++) {
      l2  += cpow(cabs(wrk->Yk[j]),2);
      err += cpow(cabs(rhs_yk),2);
      err2 += cpow(cabs(hk[3][j]),2);
    }
    err = sqrt(err/l2);
    err2 = sqrt(err2); 
    nn++;
    if ( (nn % 10)==0 ) {  
      printf("Iteration %5d. distance:  %Qe\tHeight:  %Qe\tSpeed: %.15e\tFluid Level: %Qe\n",nn, err,hmax/(2*pi), 1/sqrt(q),ml);
    }
  }
  printf("Iteration %5d. distance:  %Qe\tHeight:  %Qe\tSpeed: %.15e\tFluid Level: %Qe\n",nn, err,hmax/(2*pi), 1/sqrt(q),ml);
  ifft(wrk->Yk,wrk->Y);
  write_out(wrk->Y,"stokes_phys.dat");
  write_out(wrk->Y,"restart.dat");
  write_out(wrk->Yk,"stokes_spec.dat");

}*/





















